package com.sunday.gateway.security.authentication;

import lombok.extern.slf4j.Slf4j;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.security.core.Authentication;
import org.springframework.security.oauth2.client.InMemoryReactiveOAuth2AuthorizedClientService;
import org.springframework.security.oauth2.client.authentication.OAuth2AuthenticationToken;
import org.springframework.security.oauth2.client.web.server.AuthenticatedPrincipalServerOAuth2AuthorizedClientRepository;
import org.springframework.security.oauth2.client.web.server.ServerOAuth2AuthorizedClientRepository;
import org.springframework.security.oauth2.client.web.server.authentication.OAuth2LoginAuthenticationWebFilter;
import org.springframework.security.web.server.DefaultServerRedirectStrategy;
import org.springframework.security.web.server.ServerRedirectStrategy;
import org.springframework.security.web.server.WebFilterExchange;
import org.springframework.security.web.server.authentication.AuthenticationWebFilter;
import org.springframework.security.web.server.authentication.RedirectServerAuthenticationSuccessHandler;
import org.springframework.security.web.server.authentication.ServerAuthenticationSuccessHandler;
import org.springframework.security.web.server.savedrequest.ServerRequestCache;
import org.springframework.security.web.server.savedrequest.WebSessionServerRequestCache;
import org.springframework.util.Assert;
import org.springframework.web.server.ServerWebExchange;
import org.springframework.web.server.WebFilterChain;
import reactor.core.publisher.Mono;

import java.net.URI;
import java.util.ArrayList;
import java.util.List;

/**
 * 成功登录回调
 *
 * @see RedirectServerAuthenticationSuccessHandler
 */
@Slf4j
public class CustomRedirectServerAuthenticationSuccessHandler implements ServerAuthenticationSuccessHandler {

    private ServerOAuth2AuthorizedClientRepository authorizedClientRepository;

    private URI location = URI.create("/");

    private ServerRedirectStrategy redirectStrategy = new DefaultServerRedirectStrategy();

    private ServerRequestCache requestCache = new WebSessionServerRequestCache();

    /**
     * Creates a new instance with location of "/"
     */
    public CustomRedirectServerAuthenticationSuccessHandler(ServerOAuth2AuthorizedClientRepository authorizedClientRepository) {
        this.authorizedClientRepository = authorizedClientRepository;
    }

    /**
     * Creates a new instance with the specified location
     *
     * @param location the location to redirect if the no request is cached in
     *                 {@link #setRequestCache(ServerRequestCache)}
     */
    public CustomRedirectServerAuthenticationSuccessHandler(String location, ServerOAuth2AuthorizedClientRepository authorizedClientRepository) {
        Assert.notNull(authorizedClientRepository, "authorizedClientRepository cannot be null");
        this.location = URI.create(location);
        this.authorizedClientRepository = authorizedClientRepository;
    }

    /**
     * Sets the {@link ServerRequestCache} used to redirect to. Default is
     * {@link WebSessionServerRequestCache}.
     *
     * @param requestCache the cache to use
     */
    public void setRequestCache(ServerRequestCache requestCache) {
        Assert.notNull(requestCache, "requestCache cannot be null");
        this.requestCache = requestCache;
    }

    /**
     * Where the user is redirected to upon authentication success
     *
     * @param location the location to redirect to. The default is "/"
     */
    public void setLocation(URI location) {
        Assert.notNull(location, "location cannot be null");
        this.location = location;
    }

    /**
     * The RedirectStrategy to use.
     *
     * @param redirectStrategy the strategy to use. Default is DefaultRedirectStrategy.
     */
    public void setRedirectStrategy(ServerRedirectStrategy redirectStrategy) {
        Assert.notNull(redirectStrategy, "redirectStrategy cannot be null");
        this.redirectStrategy = redirectStrategy;
    }


    @Override
    public Mono<Void> onAuthenticationSuccess(WebFilterExchange webFilterExchange, Authentication authentication) {
        ServerWebExchange exchange = webFilterExchange.getExchange();
        ServerHttpResponse response = exchange.getResponse();

        log.info("登录成功, {} ", authentication.getPrincipal());

        if (authentication instanceof OAuth2AuthenticationToken oauth2Token) {
            String clientRegistrationId = oauth2Token.getAuthorizedClientRegistrationId();

            /**
             * 程序重启更新拉取信息
             * @see AuthenticationWebFilter#filter(ServerWebExchange, WebFilterChain)
             * @see AuthenticationWebFilter#authenticate(ServerWebExchange, WebFilterChain, Authentication)
             * @see OAuth2LoginAuthenticationWebFilter#onAuthenticationSuccess(Authentication, WebFilterExchange)
             *
             * 加载流程
             * @see AuthenticatedPrincipalServerOAuth2AuthorizedClientRepository#loadAuthorizedClient(String, Authentication, ServerWebExchange)
             * @see InMemoryReactiveOAuth2AuthorizedClientService#loadAuthorizedClient(String, String)
             */
            // 使用 flatMap 避免阻塞操作
            return authorizedClientRepository.loadAuthorizedClient(clientRegistrationId, oauth2Token, exchange)
                    .flatMap(client -> {
                        StringBuilder urlBuilder = new StringBuilder(location.toString());
                        if (client != null) {
                            List<String> paramPairs = new ArrayList<>();
                            paramPairs.add(STR."accessToken=\{client.getAccessToken().getTokenValue()}");
                            if (client.getRefreshToken() != null) {
                                paramPairs.add(STR."refreshToken=\{client.getRefreshToken().getTokenValue()}");
                            }
                            urlBuilder.append("?").append(String.join("&", paramPairs));
                        }
                        return onAuthenticationSuccess(webFilterExchange, URI.create(urlBuilder.toString()));
                    })
                    .then(
                            // 使用后清空，因使用得是jwt，所以无需保存
                            authorizedClientRepository.removeAuthorizedClient(clientRegistrationId, oauth2Token, exchange)
                    )
                    ;
        }
        return onAuthenticationSuccess(webFilterExchange, location);
    }

    public Mono<Void> onAuthenticationSuccess(WebFilterExchange webFilterExchange, URI uri) {
        ServerWebExchange exchange = webFilterExchange.getExchange();
        return this.requestCache.getRedirectUri(exchange)
                .defaultIfEmpty(uri)
                .flatMap((location) -> this.redirectStrategy.sendRedirect(exchange, location));
    }


}
